import os

from dotenv import load_dotenv
from langgraph.prebuilt import create_react_agent
from langchain.chat_models import init_chat_model
from langgraph.checkpoint.memory import InMemorySaver

from langgraph_demo.langchain_weather import get_weather
from langgraph_demo.tool_call import write_file

if __name__ == '__main__':
    checkpointer = InMemorySaver()
    tools = [get_weather, write_file]
    load_dotenv(override=True)
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    agent = create_react_agent(model=model, tools=tools, checkpointer=checkpointer)
    # agent = create_react_agent(model=model, tools=tools)
    config = {
        "configurable": {
            "thread_id": "1"
        }
    }
    response = agent.invoke({"messages": [
        {"role": "user", "content": "你好，我叫何沛，好就不见"}
    ]},config)
    print(response['messages'][-1].content)

    latest=agent.get_state(config)
    print(latest)
    config2 = {
        "configurable": {
            "thread_id": "2"
        }
    }
    response = agent.invoke({"messages": [
        {"role": "user", "content": "你好，你还记得我叫什么吗"}
    ]},config2)
    latest = agent.get_state(config)
    print(latest)
    print(response['messages'][-1].content)
